import numpy as np
import z3
import sys
import matplotlib.pyplot as plt
from mlxtend.plotting import plot_decision_regions
from lp_planner import *
from planner_utils import *
class Constraint:
    def __init__(self, func):
        self.constraint = func
    def predict(self, xp):
        res = []
        for i in xp:
            res.append(self.constraint(i[0],i[1]))
        return np.array(res).astype(int)

x,y = simple_plan(400, free_single_easy)

constraint = Constraint(free_single_easy_bool)
x_range=np.linspace(-4,4,100)
y_range=x_range
mesh_xy = np.array(np.meshgrid(x_range,y_range)).T.reshape(-1,2)
mesh_logit = constraint.predict(mesh_xy)
scatter_kwargs = {'alpha': 0.0}
fig = plot_decision_regions(mesh_xy, mesh_logit, clf=constraint, scatter_kwargs=scatter_kwargs)

for i in range(len(x)):
    fig.scatter(x, y, marker=1)
plt.savefig(sys.argv[1])
